package cn.jbolt.ai.core.fileSplit;

import cn.jbolt.ai.klb.DocDataType;
import cn.jbolt.ai.klb.entity.AiDoc;
import cn.jbolt.ai.klb.entity.AiFile;
import com.jboltai.capability.message.AIMessage;
import com.jboltai.event.EventError;
import com.jboltai.event.ai.aichat.AIChatEvent;
import com.jboltai.event.handler.AIChatEventHandler;
import com.jboltai.util.other.StringUtil;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.ArrayList;
import java.util.List;

public class AIExtractQAHandler implements AIChatEventHandler {
    private static final Logger logger = LoggerFactory.getLogger(AIExtractQAHandler.class);
    private AiFile file;

    public AIExtractQAHandler(AiFile file) {
        this.file = file;
    }


    @Override
    public void onSuccess(AIChatEvent event, AIMessage msg) {
        logger.info("从文本块中提取出QA:{}", msg.getContent());
        String content = msg.getContent();

        // 提取<ul>标签中的内容
        content = StringUtil.extractByRegex(content, "<ul>(.*?)</ul>");
        if (content != null) {
            // 提取所有<li>标签中的内容
            List<String> qaItems = StringUtil.extractAllByRegex(content, "<li>(.*?)</li>");
            List<AiDoc> docs = new ArrayList<>();

            for (String qaItem : qaItems) {
                String question = StringUtil.extractByRegex(qaItem, "<que>(.*?)</que>", 1);
                String answer = StringUtil.extractByRegex(qaItem, "<ans>(.*?)</ans>", 1);

                if (question != null && answer != null) {
                    AiDoc doc = new AiDoc();
//                    doc.setTitle(question);
                    doc.setAnswer(answer);
                    doc.setContent(question);

                    doc.setKnowledgeBaseId(file.getKnowledgeBaseId());
                    doc.setFileId(file.getId());
                    doc.setDataType(DocDataType.FILE_PIECE.getCode());
                    doc.setDocType(file.getDocType());
                    docs.add(doc);
                }
            }

            event.setAttrs("docs", docs);
        }
    }

    @Override
    public void onFail(AIChatEvent event, EventError error) {
        logger.error("AI从文本块提取QA时发生异常：{}", error.getFullMsg());
        FileSplitCache.addFailCount(file.getId());
    }

    @Override
    public void onComplete(AIChatEvent event) {

    }

}
